#!/bin/bash

. ./path.sh || exit 1;
cd ..

export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
stage=5 # start from 0 if you need to start from data_list preparation
stop_stage=5

dir=
train_config=
decode_checkpoint=
decode_checkpoint_name=
gpu_id=
data_type=raw
decode_modes="salmonn_decode"
test_data_dir="/home/work_nfs8/xlgeng/data/scp_test"
test_sets="aishell--aishell2"
. tools/parse_options.sh || exit 1;

test_sets=$(echo "$test_sets" | sed 's/--/ /g')
echo "开始打印主要变量，这些变量有命令行传入"
echo "dir=$dir"
echo "train_config=$train_config"
echo "decode_checkpoint=$decode_checkpoint"
echo "decode_checkpoint_name=$decode_checkpoint_name"
echo "gpu_id=$gpu_id"
echo "stage=$stage"
echo "stop_stage=$stop_stage"
echo "test_sets=$test_sets"
test_sets=$(echo "$test_sets" | sed 's/--/ /g')
for test_set in $test_sets; do
{
    echo "prepare test this dataset: $test_set"
}
done
wait
set -e
set -u
set -o pipefail


if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
  ctc_weight=0.5
  for test_set in $test_sets; do
  {

    data_path="none"
    if [ "$data_type" == "raw" ]; then
        data_path="$test_data_dir/$test_set/data.list"
    elif [ "$data_type" == "shard" ]; then
        data_path="$test_data_dir/$test_set/shards_list.txt"
    else
        echo "Unknown data_type: $data_type"
        # 可以在这里设置A为默认值或者退出脚本
        data_path=""
        exit 1
    fi
    echo "data_path : $data_path"
    echo "test this dataset: $test_set"
    test_dir=$dir/${decode_checkpoint_name}/${test_set}
    mkdir -p $test_dir
    # if [ -f "$test_dir/wer" ] && [ -s "$test_dir/wer" ] && [ $(wc -c < "$test_dir/text") -gt 20 ]; then
    #   echo "$test_set has been decoded! don't need to decode again!"
    #   continue
    # fi
    export CUDA_VISIBLE_DEVICES=$gpu_id
    python wenet/bin/recognize_common.py --gpu $gpu_id \
      --modes $decode_modes \
      --config $dir/train.yaml \
      --data_type $data_type \
      --test_data $data_path  \
      --checkpoint $decode_checkpoint \
      --beam_size 10 \
      --batch_size 1 \
      --penalty 0.0 \
      --result_dir $test_dir \
      --ctc_weight $ctc_weight \


    echo "$test_set has been decoded!"
    python tools/compute-wer.py --char=1 --v=1 \
      $test_data_dir/$test_set/text $test_dir/text > $test_dir/wer
  }
  done
  wait
fi
if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
  ctc_weight=0.5
  for test_set in "${test_sets[@]}"; do
  {
    echo "compute wer this dataset: $test_set"
    test_dir=$dir/test_${decode_checkpoint_name}_${test_set_suffix}/${test_set}
    python tools/compute-wer.py --char=1 --v=1 \
      $test_data_dir/$test_set/text $test_dir/text > $test_dir/wer
    echo "$test_set has been decoded!"
  }
  done
  wait

fi
